package com.fudan.run.trans;

import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function;

import com.fudan.cfg.base.MapUtils;
import com.fudan.cfg.trans.BuildNodeTransDef;
import com.fudan.run.JobRunner;
import com.fudan.run.VariableRddMap;
import com.fudan.run.ctx.annotation.TransAdapter;

import scala.Tuple2;

@TransAdapter("build_node")
public class BuildNodeTrans extends TransHandler<BuildNodeTransDef>{

	@Override
	public void trans(JobRunner jobRunner, BuildNodeTransDef transNode, VariableRddMap variableMap) {
		JavaRDD<Map<String, Object>> targetRdd = variableMap.getRdd(transNode.getTarget());
		JavaPairRDD<String, Map<String, Object>> nodePair = targetRdd.mapToPair(json->{
			String key = (String)json.get(transNode.getKey());
			if(key == null) {
				key = "null";
			}
			return Tuple2.apply(key, json);
		}).filter(pair->{
			return !pair._1.equals("null");
		});//(sn,map)
//		JavaPairRDD<String, Map<String, Object>> node = nodePair.mapToPair(pair->{
//			String nodeName = (String) pair._2().get(transNode.getNodeName());
//			Map<String,Object> nodeMap = new HashMap<>();
//			nodeMap.put(nodeName, pair._2());
//			return Tuple2.apply(pair._1(), nodeMap);
//		});
		JavaRDD<Map<String, Object>> resultRdd = nodePair.groupByKey().map(agg->{
			Iterator<Map<String, Object>> iterator = agg._2().iterator();
			Map<String,Object> result = new HashMap<>();
			String parent = transNode.getParentNode();
			List<Map<String,Object>> list = new LinkedList<>();
			String nodeName = transNode.getNodeName();
			if(!nodeName.startsWith("%")) {
				result.put(transNode.getNodeName(), list);
			}
			while(iterator.hasNext()) {
				Map<String, Object> atom = iterator.next();
				if(nodeName.startsWith("%")) {
					List<Map<String,Object>> subList = (List<Map<String, Object>>) result.get(atom.get(nodeName.substring(1)));
					if(subList == null) {
						subList = new LinkedList<>();
						result.put((String)atom.get(nodeName.substring(1)), subList);
					}
					subList.add(atom);
				}else {
					List<Map<String,Object>> subList = (List<Map<String, Object>>) result.get(transNode.getNodeName());
					if(subList == null) {
						subList = new LinkedList<>();
						result.put(transNode.getNodeName(), subList);
					}
					subList.add(atom);
				}
				List<String> extractToRoot = transNode.getExtractToRoot();
				for (String rootKey : extractToRoot) {
					if(atom.get(rootKey) != null) {
						result.put(rootKey, atom.get(rootKey));
					}
				}

			}
			Map<String,Object> ret = new HashMap<>();
			if(parent == null || "".equals(parent.trim()) || "root".equals(parent.trim())) {
				ret.putAll(result);
			}else {
				MapUtils.putValue(ret, parent, result);
			}
			return ret;
		});
		variableMap.addRdd(transNode.getName(), resultRdd);
	}

}
